package pl.jug.trojmiasto.luceneExample;

import java.io.IOException;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.FieldCache;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.function.CustomScoreProvider;
import org.apache.lucene.search.function.CustomScoreQuery;
import org.apache.lucene.search.function.ValueSourceQuery;

public class RatioBoostingQuery extends CustomScoreQuery {

	/**
	 * 
	 */
	private static final long serialVersionUID = -4663706699314423359L;

	private float lowerBound;
	private float upperBound;
	private String fieldName;

	public RatioBoostingQuery(Query subQuery, 
			float lowerBound, float upperBound, String fieldName) {
		super(subQuery);
		this.lowerBound = lowerBound;
		this.upperBound = upperBound;
		this.fieldName = fieldName;
	}

	@Override
	protected CustomScoreProvider getCustomScoreProvider(IndexReader reader)
			throws IOException {
		// TODO Auto-generated method stub
		return new RatioBooster(reader);
	}

	private class RatioBooster extends CustomScoreProvider {

		private float[] ratioValues;

		public RatioBooster(IndexReader reader) throws IOException {
			super(reader);
			ratioValues = FieldCache.DEFAULT.getFloats(reader, fieldName);
		}

		@Override
		public float customScore(int doc, float subQueryScore, float valSrcScore)
				throws IOException {
			float ratio = ratioValues[doc];
			return computeScore(subQueryScore, ratio);
		}

		private float computeScore(float subScore, float ratio) {
			return (float) (Math.sqrt((double) subScore))
					* computeDistanceFactor(ratio);
		}

		private float computeDistanceFactor(float ratio) {
			return 1 / (1 + computeDistance(ratio));
		}

		private float computeDistance(float ratio) {
			float distance;
			if (ratio <= upperBound && ratio >= lowerBound) {
				distance = 0;
			} else if (ratio >= upperBound) {
				distance = ratio - upperBound;
			} else {
				distance = lowerBound - ratio;
			}
			return distance;
		}

		private Explanation customExplain(int doc, Explanation subQueryExpl) {
			Explanation expl = new Explanation(computeScore(ratioValues[doc],
					subQueryExpl.getValue()),
					"custom score = sqrt(subQueryScore) * distanceFactor for:");

			Explanation distanceFactorExpl = new Explanation(
					computeDistanceFactor(ratioValues[doc]), "distance factor");
			expl.addDetail(distanceFactorExpl);
			expl.addDetail(subQueryExpl);
			return expl;
		}

		@Override
		public Explanation customExplain(int doc, Explanation subQueryExpl,
				Explanation valSrcExpl) throws IOException {
			return customExplain(doc, subQueryExpl);
		}

		@Override
		public Explanation customExplain(int doc, Explanation subQueryExpl,
				Explanation[] valSrcExpls) throws IOException {
			return customExplain(doc, subQueryExpl);
		}
	}

}
